from keras import Input
from keras.applications import VGG19, ResNet50V2, ResNet101V2, MobileNetV2, DenseNet121, DenseNet169, DenseNet201, \
    MobileNetV3Small, ResNet152V2, VGG16, EfficientNetV2B0, EfficientNetV2S, InceptionV3, EfficientNetV2L, \
    EfficientNetV2B3, InceptionResNetV2

from .base_layer import BaseLayer
from entity.model_describe import ModelDescribe
from log import log


# [ModelLayers.ConvBase, InceptionResNetV2, True, 200, "avg],
class ConvBaseLayer(BaseLayer):
    def transfer(self, model_list_line: list, inputs: Input, models: list = None):
        model_type = model_list_line[1]
        # 没写trainable默认False
        trainable = False if len(model_list_line) < 3 else model_list_line[2]
        pooling = None if len(model_list_line) < 5 else model_list_line[4]
        conv_base = model_type(
            include_top=False,
            pooling=pooling,
        )
        conv_base.trainable = trainable
        log.info("length:"+str(len(conv_base.layers)))
        trainable_num = 0 if len(model_list_line) < 4 else model_list_line[3]
        if trainable_num < len(conv_base.layers):
            for i in range(int(len(conv_base.layers) - trainable_num)):
                conv_base.layers[i].trainable = False
        return conv_base(inputs)

    def persistence(self, model_list_line: list, train_id: int) -> ModelDescribe:
        m = super().persistence(model_list_line, train_id)

        model_type = model_list_line[1]
        # 没写trainable默认False
        trainable = False if len(model_list_line) < 3 else model_list_line[2]
        m.var2 = trainable
        m.var1 = model_type.__name__
        m.var3 = 0 if len(model_list_line) < 4 else model_list_line[3]
        m.var4 = None if len(model_list_line) < 5 else model_list_line[4]
        return m

    def check(self, model_list_line: list, models: list = None) -> bool:
        if len(model_list_line) < 2:
            log.error("传入参数不足")
            raise RuntimeError("传入参数不足")
        model_type = model_list_line[1]
        if model_type not in [VGG19,
                              ResNet50V2,
                              ResNet101V2,
                              MobileNetV2,
                              DenseNet121,
                              DenseNet169,
                              DenseNet201,
                              MobileNetV3Small,
                              ResNet152V2,
                              VGG16,
                              EfficientNetV2B0,
                              EfficientNetV2S,
                              InceptionV3,
                              EfficientNetV2L,
                              EfficientNetV2B3,
                              EfficientNetV2B0,
                              InceptionResNetV2,
                              ]:
            log.error("未知基础模型")
            raise RuntimeError("未知基础模型")
        return True
